import random
from sqlalchemy import text
from core.generator.data_generator import DataGenerator


class MySqlDataGenerator(DataGenerator):
    def __init__(self, connector, config):
        super().__init__(connector, config)

    def generate_4_no_text(self, column_info):
        # 1. 根据字段类型直接生成数据
        if column_info['data_type'] in ["date", "datetime", "timestamp"]:
            return self.faker.date_time_this_century().strftime("%Y-%m-%d %H:%M:%S")
        elif column_info['data_type'] in ["int", "bigint", "tinyint", "smallint", "mediumint"]:
            return random.randint(1, 99999999)
        elif column_info['data_type'] in ["float", "double", "decimal"]:
            return round(random.uniform(1, 100), 2)
        elif column_info['data_type'] in ["varchar", "text", "char"]:
            return None

    from sqlalchemy import text

    def insert_row(self, connector, table, row):
        """
        插入一行数据到指定表中。
        :param connector: MySQLConnector 实例
        :param table: 表名
        :param row: 数据行（字典）
        :return: 插入成功标志和错误信息
        """
        # 构造插入 SQL
        columns = ", ".join(row.keys())
        placeholders = ", ".join([f":{key}" for key in row.keys()])
        insert_query = text(f"INSERT INTO {table} ({columns}) VALUES ({placeholders})")

        try:
            # 获取 SQLAlchemy 的会话
            with connector.get_session() as session:
                session.execute(insert_query, row)
                session.commit()
            return 1, {}
        except Exception as db_error:
            return 0, {
                "row": row,
                "error": str(db_error),
                "sql": insert_query.text
            }

